Skip to content

Add masked importance sampling for FSDP backend (#1063).#1

Merged
zijiexia merged 5 commits intozijiexia:zijie_dev_branchfrom
GuanxingLu:feature/fsdp-mis
Dec 16, 2025
Merged

Add masked importance sampling for FSDP backend (#1063).#1
zijiexia merged 5 commits intozijiexia:zijie_dev_branchfrom
GuanxingLu:feature/fsdp-mis

Conversation

@GuanxingLu
Copy link

Hi Zijie, I have my code here. Please feel free to take a look, thank you!

Summary:

  • Directly use compute_mis_weights func from megatron backend
  • Add a pytest file (tests/test_fsdp_mis.py)

@GuanxingLu
Copy link
Author

Here are the results I run with 4xH200 GPUs (using examples/train_infer_mismatch_helper/run-qwen3-4b-fsdp-mis.sh):

image

Unfortunately, the original mismatch of training and rollout engine is quite marginal, so the MIS just has no effect. Thus I add a pytest script to test the functionality.

@zijiexia zijiexia merged commit 8c8a782 into zijiexia:zijie_dev_branch Dec 16, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants